import pygame
from pygame import joystick
import json
import time
import multiprocessing as mp
import queue
import os
from consts import *
pygame.init()


def get_axis_dict(robot_axis=0, in_min=-1, in_max=1, out_min=-100, out_max=100):
    """手柄轴, 输入下限, 输入上限, 输出下限, 输出上限"""
    _dict = {}
    _dict["robot_axis"] = robot_axis
    _dict["in_min"] = in_min
    _dict["in_max"] = in_max
    _dict["out_min"] = out_min
    _dict["out_max"] = out_max
    return _dict

# 从文件中读取配置
if os.path.exists("config.json"):
    with open("config.json", "r") as f:
        CONFIG:dict = json.load(f)
    if f"axis_{1}" not in CONFIG.keys():
        # 输出轴，输入轴，输入轴范围，输出轴范围
        CONFIG[f"axis_{1}"] = get_axis_dict(0, -1, 1, -100, 100)
    if f"axis_{3}" not in CONFIG.keys():
        # 输出轴，输入轴，输入轴范围，输出轴范围
        CONFIG[f"axis_{3}"] = get_axis_dict(1, -1, 1, -100, 100)
    if f"axis_{4}" not in CONFIG.keys():
        # 输出轴，输入轴，输入轴范围，输出轴范围
        CONFIG[f"axis_{4}"] = get_axis_dict(2, -1, 1, 0,   -360)
    if f"axis_{5}" not in CONFIG.keys():
        # 输出轴，输入轴，输入轴范围，输出轴范围
        CONFIG[f"axis_{5}"] = get_axis_dict(2, -1, 1, 0,    360)
        
    if "button" not in CONFIG.keys():
        CONFIG["button"] = {"increase_speed":6,
                            "decrease_speed":7}
    if "joystick_FPS" not in CONFIG.keys():
        CONFIG["joystick_FPS"] = 60
    if "shift_cancelling_value" not in CONFIG.keys():
        CONFIG["shift_cancelling_value"] = 0.1
else:
    CONFIG = {}
    CONFIG["axis_1"] = get_axis_dict(0, -1, 1, -100, 100)
    CONFIG["axis_3"] = get_axis_dict(1, -1, 1, -100, 100)
    CONFIG["axis_4"] = get_axis_dict(2, -1, 1, 0,   -360)
    CONFIG["axis_5"] = get_axis_dict(2, -1, 1, 0,    360)
    CONFIG["button"] = [(6, INCREASE_SPEED), (7, DECREASE_SPEED)]
    CONFIG["joystick_FPS"] = 60
    CONFIG["shift_cancelling_value"] = 0.1


# 同步配置
with open("config.json", 'w') as js_file:
    json.dump(CONFIG, js_file, indent=4)

def spd_map_func(map_in:list=[-1,1], map_out:list=[-3600, 3600]):
    """将速度从区间a线性映射到区间b, 返回函数的斜率k和偏置b"""
    k = (map_out[1] - map_out[0]) / (map_in[1] - map_in[0])
    b = (map_out[0] - k*map_in[0])
    return k, b

def spd_map_func_(maps:list=[-1, 1, -3600, 3600]):
    """将速度从区间a线性映射到区间b, 返回函数的斜率k和偏置b"""
    k = (maps[3] - maps[2]) / (maps[1] - maps[0])
    b = (maps[2] - k*maps[0])
    return k, b

def axis_shift_cancelling(axis_value:float, fixValue:float)->float:
    """手柄的漂移补偿`axis_value ~ [-1, 1]`"""
    if fixValue > axis_value  > -fixValue:
        axis_value  = 0.
    elif -1. < axis_value  < -1. + fixValue:
        axis_value  = -1.
    elif 1. > axis_value  > 1. - fixValue:
        axis_value  = 1.
    return axis_value

def get_config_axis(config:dict):
    joy_axis_list   = []
    axis_parameters = []
    for i in config.keys():
        if i.startswith("axis_"):
            joy_axis_list.append(int(i.split("_")[1]))
            id    = config[i]["robot_axis"]
            in_l  = config[i]["in_min"]
            in_h  = config[i]["in_max"]
            out_l = config[i]["out_min"]
            out_h = config[i]["out_max"]
            k, b, = spd_map_func_([in_l, in_h, out_l, out_h])
            axis_parameters.append([id, k, b])
    return joy_axis_list, axis_parameters



class Joystick_Manager():
    def __init__(self) -> None:
        # self.clock = pygame.time.Clock()
        self.joystick:int = None
        self.joystick_outputs = mp.Queue(maxsize=1)
        self.parent_queue = mp.Queue()
        self.joy_list = []
        self.axis_list, self.axis_parameters = get_config_axis(CONFIG)
        self.subP = None
     
    def scan_joystick(self):
        """扫描连接到系统上的手柄"""
        joy_names = [joystick.Joystick(i).get_name() for i in range(joystick.get_count())]
        self.joy_list = joy_names
        return joy_names
    
    def stop_joystick(self):
        """关闭手柄"""
        if self.joystick is not None:
            self.parent_queue.put("stop")
            self.joystick = None
            # 清空队列
            while not self.joystick_outputs.empty():
                self.joystick_outputs.get()
            time.sleep(0.1)
            self.subP.terminate()  # 终止进程
            self.subP = None
            

    def start_joystick(self, id: int):
        """启用某个手柄"""
        self.stop_joystick()
        self.subP = mp.Process(target=self._run_mp, args=(id,))
        self.subP.name = "joystick_manager"
        self.subP.daemon = True
        self.joystick = id
        self.subP.start()


    def _run_mp(self, id):
        env = mp_joystick(id, self.joystick_outputs, self.parent_queue, 
                          self.axis_list, self.axis_parameters)
        env.run()

    def pause_joystick(self):
        if self.joystick is not None:
            print("按下")
            self.parent_queue.put("pause")

    def continue_joystick(self):
        if self.joystick is not None:
            print("放开")
            self.parent_queue.put("continue")

    def _main_loop_test(self):
        print("测试手柄功能, 将在10秒后关闭并退出子进程")
        # n = 0
        while True:
            # n += 1
            # time.sleep(0.1)
            try:
                data = self.joystick_outputs.get_nowait()
                print(data)
            except queue.Empty:
                ...
            # self.parent_queue.put(n)
            # if n > 100:
            #     self.stop_joystick()
            #     break
        self.stop_joystick()

class mp_joystick():
    """传入pygame.joystick.Joystick实例, 以及robot实例。该实例提供set_speed(int id, fload speed)方法"""
    def __init__(self, id, outconn:queue.Queue, inconn:queue.Queue, 
                 axis_list, axis_parameters, FPS=15) -> None:
        # pygame.init()
        self.clock = pygame.time.Clock()
        self.joy = joystick.Joystick(id)
        self.joy.init()
        self.FPS = FPS
        self.conn = outconn
        self.inconn = inconn
        self.isRunning = False
        self.button_num = self.joy.get_numbuttons()
        self.axis_list = axis_list
        self.axis_parameters = axis_parameters
        self.motor_axis_value = {i[0]:0. for i in axis_parameters}
        self.pause = False
    
    def run(self):
        print(f"手柄{self.joy.get_name()}进程开启")
        self.isRunning = True
        while self.isRunning:
            self.clock.tick(self.FPS)
            if not self.inconn.empty():
                msg = self.inconn.get()
                if msg == "stop":
                    self.isRunning = False
                    break
                elif msg == "pause":
                    self.pause = True
                    continue
                elif msg == "continue":
                    self.pause = False
                    continue
            if self.pause:
                continue
            for event in pygame.event.get():
                if event.type == pygame.JOYDEVICEADDED:
                    self.joy = pygame.joystick.Joystick(event.device_index)
                    print(f"Joystick {self.joy.get_instance_id()} connencted")

                elif event.type == pygame.JOYDEVICEREMOVED:
                    self.joy = None
                    print(f"Joystick {event.instance_id} disconnected")
                    self.conn.put(f"{ERROR} {JOYDISCONNECTED}")
    
                elif event.type == pygame.JOYBUTTONDOWN:
                    # 运行所有的按钮控制程序
                    self.conn.put(f"{BUTTON_PUSHED} {event.button}")
                elif event.type == pygame.JOYBUTTONUP:
                    # 运行所有的按钮控制程序
                    self.conn.put(f"{BUTTON_RELEASED} {event.button}")
                elif event.type == pygame.JOYHATMOTION:
                    # 将方向帽也视作按钮
                    if event.value[0] == 1:
                        self.conn.put(f"{BUTTON_PUSHED} {self.button_num+1}")
                    if event.value[0] == -1:
                        self.conn.put(f"{BUTTON_PUSHED} {self.button_num+2}")
                    if event.value[1] == 1:
                        self.conn.put(f"{BUTTON_PUSHED} {self.button_num+3}")
                    if event.value[1] == -1:
                        self.conn.put(f"{BUTTON_PUSHED} {self.button_num+4}")
            if self.joy is not None:
                for axis, parameters in zip(self.axis_list, self.axis_parameters):
                    joy_value = self.joy.get_axis(axis)
                    joy_value = axis_shift_cancelling(joy_value, CONFIG["shift_cancelling_value"])
                    motor_id, k, b = parameters
                    self.motor_axis_value[motor_id] += (k * joy_value) + b
                    
                msg = ""
                for motor_id in self.motor_axis_value.keys():                    
                    msg += f"{MOTOR_SPEED} {motor_id} {self.motor_axis_value[motor_id]:.2f}\n"
                    self.motor_axis_value[motor_id] = 0                    
                self.conn.put(msg)      
        print("手柄进程被终止")



if __name__ == "__main__":
    joystick_manager = Joystick_Manager()
    # joystick_manager.joystick = joystick.Joystick(0)
    print(joystick_manager.scan_joystick())
    joystick_manager.start_joystick(0)
    joystick_manager._main_loop_test()
    